# 获取至少10支股票两年的日线行情数据

import tushare as ts
import pandas as pd
from datetime import datetime, timedelta

# 设置token (需要先注册获取)
ts.set_token('1c7f85b9026518588c0d0cdac712c2d17344332c9c8cfe6bc83ee75c')
pro = ts.pro_api()

# 获取10支股票代码示例 (这里以上证50成分股为例)
stock_list = ['600519.SH', '601318.SH', '600036.SH', '601166.SH',
              '600276.SH', '600887.SH', '601398.SH', '601288.SH',
              '600000.SH', '601988.SH']

# 计算日期范围
end_date = datetime.now().strftime('%Y%m%d')
start_date = (datetime.now() - timedelta(days=365*2)).strftime('%Y%m%d')

# 获取数据
all_data = pd.DataFrame()

for stock in stock_list:
    df = pro.daily(ts_code=stock, start_date=start_date, end_date=end_date)
    df['stock'] = stock
    all_data = pd.concat([all_data, df])
    print(f"已获取 {stock} 数据，共 {len(df)} 条记录")

# 查看数据结构
print(all_data.head())
print(all_data.info())

import tushare as ts
import pandas as pd
import numpy as np
from datetime import datetime, timedelta
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.preprocessing import MinMaxScaler
from sklearn.decomposition import PCA
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import classification_report, confusion_matrix, roc_curve, auc
from ta import add_all_ta_features
from ta.utils import dropna

# 设置绘图风格和中文字体
plt.style.use('seaborn')
sns.set_palette("husl")
plt.rcParams['font.sans-serif'] = ['SimHei']  # 使用黑体
plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题

# 转换日期格式并排序
all_data['trade_date'] = pd.to_datetime(all_data['trade_date'], format='%Y%m%d')
all_data = all_data.sort_values(['stock', 'trade_date'])

# 转换数据类型
numeric_cols = ['open', 'high', 'low', 'close', 'pre_close', 'change', 'pct_chg', 'vol', 'amount']
all_data[numeric_cols] = all_data[numeric_cols].apply(pd.to_numeric, errors='coerce')

# 添加收益率标签 (未来5日收益率)
all_data['future_5d_return'] = all_data.groupby('stock')['close'].pct_change(5).shift(-5)

# 分类标签：1表示未来5日上涨，0表示下跌
all_data['label'] = (all_data['future_5d_return'] > 0).astype(int)

# 删除未来收益为NaN的行
all_data = all_data.dropna(subset=['future_5d_return'])

# 为每只股票单独计算技术指标
all_data_with_ta = pd.DataFrame()

for stock in all_data['stock'].unique():
    stock_data = all_data[all_data['stock'] == stock].copy()

    # 使用ta库计算技术指标
    stock_data = add_all_ta_features(
        stock_data,
        open="open",
        high="high",
        low="low",
        close="close",
        volume="vol",
        fillna=True
    )

    # 手动添加一些额外指标
    stock_data['ma_5'] = stock_data['close'].rolling(5).mean()
    stock_data['ma_10'] = stock_data['close'].rolling(10).mean()
    stock_data['ma_20'] = stock_data['close'].rolling(20).mean()
    stock_data['ma_60'] = stock_data['close'].rolling(60).mean()

    stock_data['boll_upper'] = stock_data['ma_20'] + 2 * stock_data['close'].rolling(20).std()
    stock_data['boll_lower'] = stock_data['ma_20'] - 2 * stock_data['close'].rolling(20).std()

    stock_data['price_ma5_ratio'] = stock_data['close'] / stock_data['ma_5']
    stock_data['price_ma20_ratio'] = stock_data['close'] / stock_data['ma_20']

    all_data_with_ta = pd.concat([all_data_with_ta, stock_data])

# 删除含有NaN的行 (由于滚动计算产生的)
all_data_with_ta = all_data_with_ta.dropna()

# 查看技术指标
print("计算的技术指标数量:", len([col for col in all_data_with_ta.columns if col not in all_data.columns]))
print("示例技术指标:", [col for col in all_data_with_ta.columns if col not in all_data.columns][:20])

# 检查缺失值
missing_values = all_data_with_ta.isnull().sum()
print("缺失值统计:\n", missing_values[missing_values > 0])


# 处理异常值 - 使用IQR方法
def handle_outliers(df, columns):
    for col in columns:
        Q1 = df[col].quantile(0.25)
        Q3 = df[col].quantile(0.75)
        IQR = Q3 - Q1
        lower_bound = Q1 - 1.5 * IQR
        upper_bound = Q3 + 1.5 * IQR

        # 将异常值替换为边界值
        df[col] = np.where(df[col] < lower_bound, lower_bound, df[col])
        df[col] = np.where(df[col] > upper_bound, upper_bound, df[col])
    return df


# 选择数值型特征处理异常值
numeric_features = all_data_with_ta.select_dtypes(include=[np.number]).columns.tolist()
all_data_cleaned = handle_outliers(all_data_with_ta.copy(), numeric_features)

# 可视化异常值处理效果
plt.figure(figsize=(12, 6))
sns.boxplot(data=all_data_with_ta[['close', 'vol', 'volume_adi', 'volatility_bbm']])
plt.title('异常值处理前', fontsize=15)
plt.xticks(rotation=45)
plt.show()

plt.figure(figsize=(12, 6))
sns.boxplot(data=all_data_cleaned[['close', 'vol', 'volume_adi', 'volatility_bbm']])
plt.title('异常值处理后', fontsize=15)
plt.xticks(rotation=45)
plt.show()

# 选择需要归一化的特征 (排除标签、日期等列)
features_to_scale = [col for col in all_data_cleaned.columns
                    if col not in ['trade_date', 'ts_code', 'stock', 'label', 'future_5d_return']]

# 初始化归一化器
scaler = MinMaxScaler()

# 归一化数据
scaled_data = scaler.fit_transform(all_data_cleaned[features_to_scale])
scaled_df = pd.DataFrame(scaled_data, columns=features_to_scale)

# 合并归一化后的数据
final_data = pd.concat([
    all_data_cleaned[['trade_date', 'ts_code', 'stock', 'label', 'future_5d_return']].reset_index(drop=True),
    scaled_df
], axis=1)

# 准备特征数据
X = final_data.drop(['trade_date', 'ts_code', 'stock', 'label', 'future_5d_return'], axis=1)
y = final_data['label']

# 执行PCA
pca = PCA()
X_pca = pca.fit_transform(X)

# 可视化解释方差比例
plt.figure(figsize=(10, 6))
plt.plot(np.cumsum(pca.explained_variance_ratio_))
plt.xlabel('主成分数量')
plt.ylabel('累计解释方差比例')
plt.title('PCA解释方差比例')
plt.grid()
plt.show()

# 选择保留95%方差的主成分
n_components = np.argmax(np.cumsum(pca.explained_variance_ratio_) >= 0.95) + 1
print(f"保留95%方差需要的主成分数量: {n_components}")

# 使用选择的主成分重新拟合PCA
pca = PCA(n_components=n_components)
X_reduced = pca.fit_transform(X)

# 可视化前两个主成分
plt.figure(figsize=(10, 6))
plt.scatter(X_reduced[:, 0], X_reduced[:, 1], c=y, alpha=0.5, cmap='viridis')
plt.xlabel('第一主成分')
plt.ylabel('第二主成分')
plt.title('PCA降维可视化 (前两个主成分)')
plt.colorbar()
plt.show()

import plotly.express as px
import plotly.graph_objects as go
from plotly.subplots import make_subplots

# 1. 选择目标股票和6个属性
target_stock = '600519.SH'  # 贵州茅台
features = [
    'close',          # 收盘价
    'vol',            # 成交量
    'momentum_rsi',   # RSI(14)
    'trend_macd',     # MACD
    'volatility_bbh', # 布林带上轨
    'volume_adi'      # 累积分布指标
]

# 2. 提取该股票数据
stock_data = final_data[final_data['stock'] == target_stock].copy()
stock_data['date'] = stock_data['trade_date'].dt.strftime('%Y-%m-%d')  # 格式化日期

# 3. 创建可交互图表
fig = make_subplots(specs=[[{"secondary_y": True}]])

# 添加收盘价（主坐标轴）
fig.add_trace(
    go.Scatter(
        x=stock_data['date'],
        y=stock_data['close'],
        name="收盘价",
        line=dict(color='royalblue', width=2),
        hovertemplate="<b>日期</b>: %{x}<br><b>收盘价</b>: %{y:.2f}<extra></extra>"
    ),
    secondary_y=False
)

# 添加其他指标（次坐标轴）
colors = px.colors.qualitative.Plotly[1:]  # 避免与收盘价颜色重复
for i, feature in enumerate(features[1:]):  # 跳过已添加的close
    fig.add_trace(
        go.Scatter(
            x=stock_data['date'],
            y=stock_data[feature],
            name=feature,
            line=dict(color=colors[i], width=1.5, dash='dot'),
            visible="legendonly",  # 默认隐藏，可通过图例切换
            hovertemplate=f"<b>日期</b>: %{{x}}<br><b>{{{feature}}}</b>: %{{y:.2f}}<extra></extra>"
        ),
        secondary_y=True
    )

# 4. 设置图表样式
fig.update_layout(
    title=f'<b>{target_stock} 六维指标时序分析</b>',
    xaxis=dict(
        title='日期',
        rangeslider=dict(visible=True),  # 添加范围滑块
        type='date'  # 确保日期格式正确解析
    ),
    yaxis=dict(
        title='收盘价',
        titlefont=dict(color='royalblue'),
        tickfont=dict(color='royalblue')
    ),
    yaxis2=dict(
        title='技术指标值',
        titlefont=dict(color='grey'),
        tickfont=dict(color='grey'),
        overlaying='y',
        side='right'
    ),
    hovermode='x unified',  # 显示同一x值所有y值
    legend=dict(
        orientation="h",
        yanchor="bottom",
        y=1.02,
        xanchor="right",
        x=1
    ),
    template='plotly_white',
    width=1000,
    height=600,
    margin=dict(autoexpand=True)
)

# 5. 添加指标筛选按钮
buttons = []
for i, feature in enumerate(features):
    visibility = [False] * len(features)
    visibility[i] = True
    buttons.append(
        dict(
            label=feature,
            method="update",
            args=[{"visible": [True] + visibility},  # 始终显示收盘价
                 {"title": f"<b>{target_stock} - {feature}指标</b>"}]
        )
    )

fig.update_layout(
    updatemenus=[{
        "buttons": buttons,
        "direction": "down",
        "showactive": True,
        "x": 1.05,
        "y": 1.15
    }]
)

# 6. 显示图表
fig.show()

# 检查标签分布
label_dist = final_data['label'].value_counts()
print("原始标签分布:\n", label_dist)

# 方法1：手动欠采样
minority_class = final_data[final_data['label'] == 1]
majority_class = final_data[final_data['label'] == 0]

# majority_downsampled = majority_class.sample(n=len(minority_class), random_state=42)
# balanced_data = pd.concat([majority_downsampled, minority_class])

# 方法2：或者使用过采样
minority_upsampled = minority_class.sample(n=len(majority_class), replace=True, random_state=42)
balanced_data = pd.concat([majority_class, minority_upsampled])

print("\n平衡后的标签分布:\n", balanced_data['label'].value_counts())

# 可视化对比
plt.figure(figsize=(12, 4))
plt.subplot(1, 2, 1)
label_dist.plot(kind='bar')
plt.title('原始数据分布')

plt.subplot(1, 2, 2)
balanced_data['label'].value_counts().plot(kind='bar')
plt.title('平衡后数据分布')

plt.tight_layout()
plt.show()

# 准备特征和标签
X = balanced_data.drop(['trade_date', 'ts_code', 'stock', 'label', 'future_5d_return'], axis=1)
y = balanced_data['label']

# 分割数据集
X_train, X_test, y_train, y_test = train_test_split(
    X, y,
    test_size=0.3,
    random_state=42,
    stratify=y  # 保持分层抽样
)

print(f"训练集大小: {X_train.shape}, 测试集大小: {X_test.shape}")

# 初始化随机森林分类器
rf = RandomForestClassifier(
    n_estimators=100,
    max_depth=10,
    min_samples_split=5,
    random_state=42,
    class_weight='balanced'
)

# 训练模型
rf.fit(X_train, y_train)

# 预测测试集
y_pred = rf.predict(X_test)
y_pred_proba = rf.predict_proba(X_test)[:, 1]  # 正类的概率

from sklearn.metrics import accuracy_score

# 预测测试集
y_pred = rf.predict(X_test)
y_pred_proba = rf.predict_proba(X_test)[:, 1]  # 正类的概率

# 计算准确率
accuracy = accuracy_score(y_test, y_pred)
print(f"模型准确率: {accuracy:.4f}")

# 1. 分类报告（已包含准确率）
print("\n分类报告:\n", classification_report(y_test, y_pred))

# 2. 混淆矩阵（增加准确率标注）
cm = confusion_matrix(y_test, y_pred)
plt.figure(figsize=(6, 6))
sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', cbar=False)
plt.xlabel('预测标签')
plt.ylabel('真实标签')
plt.title(f'混淆矩阵 (准确率: {accuracy:.2%})')  # 添加准确率到标题
plt.show()

# 3. ROC曲线（增加准确率标注）
fpr, tpr, thresholds = roc_curve(y_test, y_pred_proba)
roc_auc = auc(fpr, tpr)

plt.figure(figsize=(8, 6))
plt.plot(fpr, tpr, color='darkorange', lw=2,
         label=f'ROC曲线 (AUC = {roc_auc:.2f}, 准确率 = {accuracy:.2f})')  # 添加准确率
plt.plot([0, 1], [0, 1], color='navy', lw=2, linestyle='--')
plt.xlim([0.0, 1.0])
plt.ylim([0.0, 1.05])
plt.xlabel('假正率(FPR)')
plt.ylabel('真正率(TPR)')
plt.title('接收者操作特征(ROC)曲线')
plt.legend(loc="lower right")
plt.show()

# 4. 特征重要性（保持原样）
feature_importance = pd.DataFrame({
    'feature': X.columns,
    'importance': rf.feature_importances_
}).sort_values('importance', ascending=False)

plt.figure(figsize=(12, 8))
sns.barplot(x='importance', y='feature', data=feature_importance.head(20))
plt.title(f'Top 20 重要特征 (模型准确率: {accuracy:.2%})')  # 添加准确率到标题
plt.show()
